import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import numpy as np
from sklearn.model_selection import train_test_split

class CNNModel(nn.Module):
    def __init__(self, input_channels, num_classes):
        super(CNNModel, self).__init__()
        
        # 特征提取层
        self.features = nn.Sequential(
            # 第一个卷积块
            nn.Conv1d(input_channels, 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # 第二个卷积块
            nn.Conv1d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # 第三个卷积块
            nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # 第四个卷积块
            nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )
        
        # 计算全连接层的输入特征数 (1200 -> 600 -> 300 -> 150 -> 75)
        self._to_linear = 128 * 75
        
        # 分类层
        self.classifier = nn.Sequential(
            nn.Linear(self._to_linear, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x, return_features=False):
        # 确保输入维度正确
        if x.dim() == 2:
            x = x.unsqueeze(1)  # 添加通道维度
        
        # 特征提取
        features = self.features(x)
        
        # 展平
        features = features.view(features.size(0), -1)
        
        # 通过第一个全连接层
        features = self.classifier[0](features)

        if return_features:
            return features
        
        # 继续通过剩余的分类层
        x = self.classifier[1:](features)
        return x

def get_model(num_classes, device):
    model = CNNModel(input_channels=1, num_classes=num_classes).to(device)
    return model

def load_and_preprocess_data(use_gan=False, gan_type="vanilla", use_sim=False):
    """加载已经划分好的训练和测试数据"""
    
    # 加载测试集 - 始终使用原始测试数据
    test_data_path = os.path.join("./cwru_prepro", "test")
    test_files = [f for f in os.listdir(test_data_path) if f.endswith('.npy')]
    
    X_test = []
    y_test = []
    class_names = []
    
    for class_idx, class_file in enumerate(test_files):
        class_name = class_file.replace('.npy', '')
        data = np.load(os.path.join(test_data_path, class_file))
        X_test.append(data)
        y_test.append(np.full(len(data), class_idx))
        class_names.append(class_name)
    
    X_test = np.concatenate(X_test)
    y_test = np.concatenate(y_test)
    
    # 根据模式加载训练数据
    if use_sim:
        # 使用仿真数据作为训练集
        X_train, y_train = load_simulated_data(class_names)
    elif use_gan:
        # 使用GAN生成数据作为训练集
        X_train, y_train = load_gan_data(gan_type, class_names)
    else:
        # 使用原始训练数据
        X_train, y_train = load_original_train_data(class_names)
    
    return X_train, X_test, y_train, y_test

def load_original_train_data(class_names):
    """加载原始训练数据"""
    train_data_path = os.path.join("./cwru_prepro", "train")
    X_train = []
    y_train = []
    
    for class_idx, class_name in enumerate(class_names):
        data = np.load(os.path.join(train_data_path, f"{class_name}.npy"))
        X_train.append(data)
        y_train.append(np.full(len(data), class_idx))
    
    return np.concatenate(X_train), np.concatenate(y_train)

def load_simulated_data(class_names):
    """加载仿真训练数据"""
    sim_data_path = "./cwru_sim"
    X_train = []
    y_train = []
    
    for class_idx, class_name in enumerate(class_names):
        data = np.load(os.path.join(sim_data_path, f"{class_name}.npy"))
        X_train.append(data)
        y_train.append(np.full(len(data), class_idx))
    
    return np.concatenate(X_train), np.concatenate(y_train)

def load_gan_data(gan_type, class_names):
    """加载GAN生成的训练数据"""
    gan_data_path = os.path.join("./cwru_gan", gan_type)
    X_train = []
    y_train = []
    
    for class_idx, class_name in enumerate(class_names):
        data = np.load(os.path.join(gan_data_path, f"{class_name}.npy"))
        X_train.append(data)
        y_train.append(np.full(len(data), class_idx))
    
    return np.concatenate(X_train), np.concatenate(y_train)